from flask import Flask, render_template, send_from_directory,request
import pymongo
import jieba
import re
from collections import Counter
import json
from fuzzywuzzy import fuzz


client = pymongo.MongoClient(host='127.0.0.1', port=27017)
# db 表示数据库名称，username 用户名 password 密码
#client.db.authenticate(useranme, password)
# 建立数据库连接
db = client.CNNVD_vul # or db = client['test']
collection = db.vuldb # or collection = db['students']



#根据漏洞名称或者名称的一部分查询漏洞的cnnvd编号（此步骤是在模糊匹配找到相应漏洞名称之后获得完整的漏洞信息 ）
def search_name_cnnvdcode(string,num):
    #查询，并且返回特定的属性，id默认返回，可以加上"_id":0
    res = collection.find({"vulName":string}, {"vulName":1,"cnnvdCode": 1,"cveCode":1,"vulTypeName":1,"affectedVendor":1, "vulDesc":1, "referUrl":1,"createTime":1,"updateTime":1,"isOfficial":1,"hazardLevel":1})
    allldata=[]

    #查询到的结果输出（第一个记录详细信息，后续重名的只记录CNNVD、CVE编号以及漏洞严重等级）
    for index,doc in enumerate(res):
        ldata=[]
        if index == 0:
            # 对每个文档执行操作
            vulName = doc["vulName"]
            cnnvdCode = doc["cnnvdCode"]
            cveCode = doc["cveCode"]
            vulTypeName = doc["vulTypeName"]
            affectedVendor = doc["affectedVendor"]
            vulDesc = doc["vulDesc"]
            referUrl = doc["referUrl"]
            createTime = doc["createTime"]
            updateTime = doc["updateTime"]
            isOfficial = doc["isOfficial"]
            hazardLevel = doc["hazardLevel"]
            
            # 打印文档中的某些字段
            ldata.append(vulName)
            ldata.append(cnnvdCode)
            ldata.append(cveCode)
            ldata.append(vulTypeName)
            ldata.append(affectedVendor)
            ldata.append(vulDesc)
            ldata.append(referUrl)
            ldata.append(createTime)
            ldata.append(updateTime)
            io=["是","否"]
            ldata.append(io[int(isOfficial)])
            hl=["未定级","超危","高危","中危","低危"]
            if hazardLevel=="null":
                ldata.append(hl[0])
            else:
                ldata.append(hl[int(hazardLevel)])
            ldata.append(num)
            allldata.append(ldata)
        else:
            cnnvdCode = doc["cnnvdCode"]
            cveCode = doc["cveCode"]

            ldata.append(cnnvdCode)
            ldata.append(cveCode)

            hl=["未定级","超危","高危","中危","低危"]
            if hazardLevel=="null":
                ldata.append(hl[0])
            else:
                ldata.append(hl[int(hazardLevel)])
            allldata.append(ldata)

    return allldata

#jieba分词，对输入的字符串进行分词，返回一个列表
def jieba_cut(string):
    seg_list = jieba.cut(string,cut_all=False)
    actual_list = [word.strip() for word in seg_list if word.strip()]   
    return actual_list


#纠错算法，依据corpus.txt语料库文件对输入的字符串进行纠错，返回正确的单词
def correct(string):
    # 定义一个函数，用于将文本中的单词提取出来
    def words(text):
        return re.findall(r'\S+', text)
    with open('corpus.txt', 'r', encoding="utf-8") as corpus_file:
        corpus = words(corpus_file.read())
    
    # 定义一个函数，用于将单词和其出现频率组成元组
    word_counter = Counter(corpus)
    word = string
    matches = []
    # 如果单词是语料库中单词的子串,则减去两者相差的单词长度（如果完全相同，则表示在忽略大小写的情况下相似度为100%）
    for candidate in word_counter.items():
        if word.lower() in candidate[0].lower() or candidate[0].lower() in word.lower():
            if word.lower() == candidate[0].lower():
                if word == candidate[0]:
                    similarity = 120
                else:
                    similarity = 110
            else:
                similarity = 100 - abs(len(candidate[0].lower())-len(word.lower()))
        # 如果单词不是语料库中单词的子串,则使用Levenshtein距离计算相似度值
        else:
            similarity = fuzz.ratio(word.lower(), candidate[0].lower())
        matches.append([candidate[0], similarity, candidate[1]])
    # 按照相似度降序排列
    sorted_matches = sorted(matches, key=lambda x: x[1], reverse=True)
    sorted_matches = sorted_matches[:20]
    min_val = min(sorted_matches, key=lambda x: x[2])  # 找到第3维度的值最小的元素
    max_val = max(sorted_matches, key=lambda x: x[2])  # 找到第3维度的值最大的元素

    # 将出现频率归一化，使频率值在0~5之间，并与相似度值加权计算，并按照加权和降序排列
    sorted_matches = sorted(sorted_matches, key=lambda x: x[1]+(x[2]-min_val[2])*5/(max_val[2]-min_val[2]), reverse=True)

    best_match = sorted_matches[0][0]  # 获取相似度最高的匹配项
    # print(f"Misspelled word: {word}")
    # print(f"Corrected word: {best_match}")
    return best_match


#合并列表中重复的元素，并统计每个元素的个数
def merge_and_count_elements(lst):
    counter = Counter(lst)
    merged_list = []
    merged_count = []
    
    for element, count in counter.items():
        merged_list.append(element)
        merged_count.append(count)
    
    return merged_list, merged_count


#jieba搜索，对输入的分词列表，按照一个分词一个分词地查找(
#这里的搜索并未去重出于两点原因：
#1.去重之后再统计同名CNNVD-CVE编号的漏洞需要重新从整个数据库中搜索计算。
#2.未去重可以突出漏洞名称的频数优势，其刚好符合搜索算法的一个基准点：基于频率的统计搜索。
def jieba_search(word_list):

    # 测试字符是否为特殊字符，因为有些特殊字符如“*”、“\”.....在pymongo中无法作为索引
    def is_special_character(char):
        escaped_char = re.escape(char)
        return escaped_char != char
    
    chr1="漏洞"
    for i in word_list:
        if is_special_character(i):
            continue
        else:
            chr1=i
            break


    print(chr1)
    # 按照第一个分词进行查找
    query = {
    "vulName": {
        "$regex": f".*{chr1}.*",
        "$options": "i"  # 不区分大小写
    }
}
    results = collection.find(query,{"vulName":1, "_id": 0})
    results_list = [result["vulName"] for result in results]
    results_list = list(set(results_list))

    #如果首次查找的结果过多（>100），且存在后续分词，则对后续分词进行查找
    if len(results_list) > 100 and len(word_list) > 1:
        for word in word_list:
            new_list=[]
            for l in results_list:
                if word.lower() in l.lower():
                    new_list.append(l)

            #当一个分词的查找结果等于0时，则停止查找
            if len(new_list) == 0:
                break
            else:
                results_list = new_list

        #将查询结果中相同的元素合并并统计数量
        results1, results2=merge_and_count_elements(results_list)
        return results1, results2

    else:
        #将查询结果中相同的元素合并并统计数量
        results1, results2=merge_and_count_elements(results_list)
        return results1, results2

    


# 计算levenstain编辑距离
def levenshtein_distance(s1, s2):
    m = len(s1)
    n = len(s2)

    # 创建一个 (m+1) x (n+1) 的二维列表
    dp = [[0] * (n + 1) for _ in range(m + 1)]

    # 初始化第一行和第一列
    for i in range(m + 1):
        dp[i][0] = i
    for j in range(n + 1):
        dp[0][j] = j

    # 动态规划计算编辑距离
    for i in range(1, m + 1):
        for j in range(1, n + 1):
            if s1[i - 1] == s2[j - 1]:
                dp[i][j] = dp[i - 1][j - 1]
            else:
                dp[i][j] = min(dp[i - 1][j] + 1,  # 删除操作
                               dp[i][j - 1] + 1,  # 插入操作
                               dp[i - 1][j - 1] + 1  # 替换操作
                               )

    return dp[m][n]


# 两个字符串的有序公共子串,如果有序公共子串的长度大于2，表明两个字符串比较接近，则返回子串长度，否则返回false
def find_common_ordered_subset(a, b):
    common_subset = ""
    i = 0
    for char_a in a:
        for char_b in b[i:]:
            if char_a == char_b:
                common_subset += char_a
                i = b.index(char_b) + 1
                break
    if len(common_subset) >= 2:
        return len(common_subset)
    else:
        return False




# 压缩排序，根据列表的第三个值的顺序将三个列表同时排序，默认从小到大，也可以倒序sorted(combined, reverse=True)
def sort_lists(list1, list2, list3):
    combined = zip(list3, list1, list2)
    sorted_combined = sorted(combined)
    sorted_list3, sorted_list1, sorted_list2 = zip(*sorted_combined)
    return list(sorted_list1), list(sorted_list2), list(sorted_list3)




def save_json(lst):
    vullist=[]
    for vul in lst:
        #print(vul)
        vuldict={}
        for index, item in enumerate(vul):
            if index == 0:
                print(item)
                vuldict["hazardLevel"]=item[10]
                vuldict["vulName"]=item[0]
                vuldict["cnnvd/cveCode"]=item[1]+"\t"+item[2]
                vuldict["vulTypeName"]=item[3]
                vuldict["affectedVendor"]=item[4]
                vuldict["vulDesc"]=item[5]                
                vuldict["referUrl"]=item[6]
            else:
                vuldict["cnnvd/cveCode"]=vuldict["cnnvd/cveCode"]+"  -/-  "+item[0]+"\t"+item[1]+"   "+item[2]

        vullist.append(vuldict)
    
    data={
            "status": 200,
            "message": "",
            "total": 8,
            "rows": {
                "item": vullist
            }
    }
    # 将数据重新编码为带缩进的JSON格式
    json_data2 = json.dumps(data, ensure_ascii=False, indent=4)

    # 保存重新编码后的数据到文件
    with open(r"static\json\2\table\demo3.json", "w",encoding="utf-8") as file2_formatted:
        file2_formatted.write(json_data2)


    file2_formatted.close()






# 执行搜索算法
def perform_search(query):
    # 在这里执行实际的搜索操作，返回搜索结果
    query=jieba_cut(query)
    print("纠错前：",query)
    # 对各个分词进行纠错
    for i in range(len(query)):
        query[i]=correct(query[i])
    print("纠错后：",query)

    # 按照jieba分词进行搜索，返回漏洞的名称及该名称漏洞的数目
    result1,result2=jieba_search(query)
    #用于存放上述返回的漏洞名称与输入的漏洞名称的匹配度
    result3=[]


    #抽象化query，将各个单词变为123的形式，将其组合为一个由数字组成的字符串。
    #计算query对应的字典
    querydict={}
    for i in range(len(query)):
        querydict[query[i]]=chr(i+65)
    #计算query对应的字符串
    querystr = ""
    for item in query:
        querystr += str(querydict[item])


    #计算result1中各个result对应的字符串，并计算其与querystr的编辑距离
    for i in result1:
        result=jieba_cut(i)

        #print(result)

        #抽象化result，将各个单词变为12345的形式，将其组合为一个由数字组成的字符串。
        #计算result对应的字典
        count=len(query)    
        resultdict={}
        for i in result:
            if i in querydict:
                resultdict[i]=querydict[i]
            else:
                resultdict[i]=chr(count+65)
                count+=1
        #计算result对应的字符串
        resultstr = ""
        for item in result:
            resultstr += str(resultdict[item])

        #顺序优先原则，如何检索到的漏洞的关键词的顺序与输入的关键词的顺序一致，则判别为相似度较高
        #计算querystr和resultstr的编辑距离,如果querystr是resultstr的有序子集，则编辑距离为0
        if find_common_ordered_subset(querystr, resultstr):
            #print("有公共子串:",len(result)-len(query))
            result3.append(abs(len(resultstr)-len(querystr))-find_common_ordered_subset(querystr, resultstr))
        else:
            result3.append(levenshtein_distance(querystr, resultstr))
    #如果result1为空，则未找到匹配项，将改提示信息返回到前端
    if len(result1)==0:
        print("没有搜索到结果")
        searchall=[[['无匹配项', '--', '--', '其他', '个人开发者', '无匹配项', '无匹配项', '2022-09-21', '2023-09-24', '否', '无匹配项', 1]]]
        return searchall
    #查询结果非空，对result1（漏洞名称）, result2（漏洞数目）, result3（相似度）按照result3的大小同时进行由小到大排序（小数在前表示相似度高），将结果返回到前端
    else:
        sorted_list1, sorted_list2, sorted_list3 = sort_lists(result1, result2, result3)
        
        if len(sorted_list3)>15:
            run=15
        else:
            run=len(sorted_list3)
        
        searchall=[]
        for i in range(run):
            ldata=search_name_cnnvdcode(sorted_list1[i], sorted_list2[i])
            searchall.append(ldata)
        return searchall
    
    
    


app = Flask(__name__)

@app.route('/')
def index():
    return render_template('index.html')

# 路由用于提供静态文件
@app.route('/static/json/2/table/demo3.json')
def serve_json():
    return send_from_directory('static', 'json/2/table/demo3.json')


# 接收输入的搜索请求并处理
@app.route('/search', methods=['GET'])
def handle_search():
    query = request.args.get('query')
    print(query)

    search_all=perform_search(query)
    save_json(search_all)

    return render_template('index.html')

@app.route('/免责声明')
def templates_statement():
    return render_template('免责声明.html')

if __name__ == '__main__':

    data={
            "status": 200,
            "message": "",
            "total": 8,
            "rows": {
                "item": [{},{},{},{},{},{},{},{}]
            }
    }
    # 将数据重新编码为带缩进的JSON格式
    json_data2 = json.dumps(data, ensure_ascii=False, indent=4)
    # 打开 JSON 文件
    with open(r"static\json\2\table\demo3.json", "w", encoding="utf-8") as file:
        file.write(json_data2)
    # 关闭文件
    file.close()

    app.run()











#   mongod --dbpath=..\data\db
